using System;
using System.Collections.Generic;
using System.Linq;
using Rhizome.Contract;
using Rhizome.Contract.Visitor;

namespace Rhizome.Model.Visitor
{
	public static class ElementEnumerator
	{
		/// <summary>
		/// 	Gets root & all elements with children
		/// </summary>
		/// <typeparam name = "TElement"></typeparam>
		/// <param name = "root"></param>
		/// <returns></returns>
		public static IEnumerable<TElement> GetAllElements<TElement>(this object root, bool includeSubAutomata)
			where TElement : class
		{
			if (root == null) throw new ArgumentNullException("root");

			var e = root as TElement;
			if (e != null) yield return e;

			var elementContainer = root as IElementContainer;
			if (elementContainer == null) yield break;

			var children =
				from element in elementContainer.GetElements()
				from subChildren in element.GetAllElements<IElement>()
				where includeSubAutomata || element is IAutomaton == false
				select subChildren;

			foreach (var item in children.OfType<TElement>())
				yield return item;
		}

		public static IEnumerable<TElement> GetAllElements<TElement>(this object root) where TElement : class
		{
			return GetAllElements<TElement>(root, true);
		}
	}
}